
#include "RainbowTableGenerator.h"
#include "ChainWalkContext.h"
#include <iostream>
#include <sstream>
#include <time.h>

char alpha[]						= "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
char alphaspace[]					= "ABCDEFGHIJKLMNOPQRSTUVWXYZ ";
char alphanumeric[]					= "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
char alphanumericspace[]			= "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 ";
char alphanumericsymbol14[]			= "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()-_+=";
char alphanumericsymbol14space[]	= "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()-_+= ";
char alphanumericsymbol32[]			= "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()-_+=~`[]{}|\\:;\"'<>,.?/";
char alphanumericsymbol32space[]	= "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()-_+=~`[]{}|\\:;\"'<>,.?/ ";

char numeric[]							= "0123456789";
char numericspace[]						= "0123456789 ";
char loweralpha[]						= "abcdefghijklmnopqrstuvwxyz";
char loweralphaspace[]					= "abcdefghijklmnopqrstuvwxyz ";
char loweralphanumeric[]				= "abcdefghijklmnopqrstuvwxyz0123456789";
char loweralphanumericspace[]			= "abcdefghijklmnopqrstuvwxyz0123456789 ";
char loweralphanumericsymbol14[]		= "abcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*()-_+=\"";
char loweralphanumericsymbol14space[]	= "abcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*()-_+=\" ";
char loweralphanumericsymbol32[] 		= "abcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*()-_+=~`[]{}|\\:;\"'<>,.?/";
char loweralphanumericsymbol32space[]	= "abcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*()-_+=~`[]{}|\\:;\"'<>,.?/ ";

char mixalpha[]							= "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
char mixalphaspace[]					= "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ ";
char mixalphanumeric[]					= "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
char mixalphanumericspace[]				= "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 ";
char mixalphanumericsymbol14[]			= "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()-_+=";
char mixalphanumericsymbol14space[]		= "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()-_+= ";
char mixalphanumericsymbol32[]			= "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()-_+=~`[]{}|\\:;\"'<>,.?/";
char mixalphanumericsymbol32space[]		= "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()-_+=~`[]{}|\\:;\"'<>,.?/ ";

CRainbowTableGenerator::CRainbowTableGenerator()
{
//	m_nCalculationSpeed = 0;
	m_nCurrentCalculatedChains = 0;

}

CRainbowTableGenerator::~CRainbowTableGenerator(void)
{
}
char* CRainbowTableGenerator::GetCharset(std::string sCharsetName, char *pCharset)
{
	if(sCharsetName == "alpha") {		
		pCharset = alpha;
	} else if(sCharsetName == "alpha-space") {		
		pCharset = alphaspace;
	} else if(sCharsetName == "alpha-numeric") {		
		pCharset = alphanumeric;
	} else if(sCharsetName == "alpha-numeric-space") {		
		pCharset = alphanumericspace;
	} else if(sCharsetName == "alpha-numeric-symbol14") {		
		pCharset = alphanumericsymbol14;
	} else if(sCharsetName == "alpha-numeric-symbol14-space") {		
		pCharset = alphanumericsymbol14space;
	} else if(sCharsetName == "alpha-numeric-symbol32") {		
		pCharset = alphanumericsymbol32;
	} else if(sCharsetName == "alpha-numeric-symbol32-space") {		
		pCharset = alphanumericsymbol32space;
	} else if(sCharsetName == "alpha-numeric-all-space") {		
		pCharset = alphanumericsymbol32space;
	} else if(sCharsetName == "numeric") {		
		pCharset = numeric;
	} else if(sCharsetName == "numeric-space") {		
		pCharset = numericspace;
	} else if(sCharsetName == "loweralpha") {		
		pCharset = loweralpha;
	} else if(sCharsetName == "loweralpha-space") {		
		pCharset = loweralphaspace;
	} else if(sCharsetName == "loweralpha-numeric") {		
		pCharset = loweralphanumeric;
	} else if(sCharsetName == "loweralpha-numeric-space") {		
		pCharset = loweralphanumericspace;
	} else if(sCharsetName == "loweralpha-numeric-symbol14") {		
		pCharset = loweralphanumericsymbol14;
	} else if(sCharsetName == "loweralpha-numeric-symbol14-space") {		
		pCharset = loweralphanumericsymbol14space;
	} else if(sCharsetName == "loweralpha-numeric-symbol32") {		
		pCharset = loweralphanumericsymbol32;
	} else if(sCharsetName == "loweralpha-numeric-symbol32-space") {		
		pCharset = loweralphanumericsymbol32space;
	} else if(sCharsetName == "loweralpha-numeric-all-space") {		
		pCharset = loweralphanumericsymbol32space;
	} else if(sCharsetName == "mixalpha") {		
		pCharset = mixalpha;
	} else if(sCharsetName == "mixalpha-space") {		
		pCharset = mixalphaspace;
	} else if(sCharsetName == "mixalpha-numeric") {		
		pCharset = mixalphanumeric;
	} else if(sCharsetName == "mixalpha-numeric-space") {		
		pCharset = mixalphanumericspace;
	} else if(sCharsetName == "mixalpha-numeric-symbol14") {		
		pCharset = mixalphanumericsymbol14;
	} else if(sCharsetName == "mixalpha-numeric-symbol14-space") {		
		pCharset = mixalphanumericsymbol14space;
	} else if(sCharsetName == "mixalpha-numeric-symbol32") {		
		pCharset = mixalphanumericsymbol32;
	} else if(sCharsetName == "mixalpha-numeric-symbol32-space") {		
		pCharset = mixalphanumericsymbol32space;
	} else if(sCharsetName == "mixalpha-numeric-all-space") {		
		pCharset = mixalphanumericsymbol32space;
	}
	return pCharset;
}
int CRainbowTableGenerator::CalculateTable(std::string sFilename, int nRainbowChainCount, std::string sHashRoutineName, std::string sCharsetName, int nPlainLenMin, int nPlainLenMax, int nRainbowTableIndex, int nRainbowChainLen, std::string sSalt)
{

//	std::fstream test("c:\\distrrtgen.log");
//	std::streambuf *old = std::cout.rdbuf(test.rdbuf());
	// Setup CChainWalkContext
	if (!CChainWalkContext::SetHashRoutine(sHashRoutineName))
	{
		std::cout << "hash routine " << sHashRoutineName << " not supported" << std::endl;
		return 1;
	}
	char *pCharset;	
	pCharset = GetCharset(sCharsetName, pCharset);

	if (!CChainWalkContext::SetPlainCharset(pCharset, strlen(pCharset), nPlainLenMin, nPlainLenMax))
		return 2;
	if (!CChainWalkContext::SetRainbowTableIndex(nRainbowTableIndex))
	{
		std::cout << "invalid rainbow table index " << nRainbowTableIndex << std::endl;
		return 3;
	}
//	std::cout.rdbuf(old);
	if(sHashRoutineName == "mscache")// || sHashRoutineName == "lmchall" || sHashRoutineName == "halflmchall")
	{
		// Convert username to unicode
		const char *szSalt = sSalt.c_str();
		int salt_length = strlen(szSalt);
		unsigned char cur_salt[256];
		for (int i=0; i<salt_length; i++)
		{
			cur_salt[i*2] = szSalt[i];
			cur_salt[i*2+1] = 0x00;
		}
		CChainWalkContext::SetSalt(cur_salt, salt_length*2);
	}
	else if(sHashRoutineName == "halflmchall")
	{
//		char salt[] = "1122334455667788";
	//	__int64 salt = 0x8877665544332211;
		//memcpy(salt, , 8);
	//	CChainWalkContext::SetSalt((unsigned char*)&salt, 8);
	}
	else if(sHashRoutineName == "oracle")
	{
		CChainWalkContext::SetSalt((unsigned char *)sSalt.c_str(), sSalt.length());
	}
	CChainWalkContext::Dump();	

	// FileName
//	std::stringstream szFilename;
//	szFilename << partid << ".rt";
//	std::ofstream Partfile;
//	Partfile.open(sFilename.c_str(), std::ios::out | std::ios::binary | std::ios::app);
	
	// Open file
	fclose(fopen(sFilename.c_str(), "a"));
	FILE* file = fopen(sFilename.c_str(), "r+b");
	
	if (file == NULL)
	{
		std::cout << "failed to create " << sFilename << std::endl;
		return 4;
	}
	
	
	// Check existing chains
	unsigned int nDataLen = (unsigned int)GetFileLen(file);
	fseek(file, nDataLen, SEEK_SET);
	unsigned int nFileLen;
	/*
	long begin,end;
	
	begin = Partfile.tellp();
	Partfile.seekp(0, std::ios::end);
	end = Partfile.tellp();
	unsigned int nDataLen = end - begin;
	*/
	nDataLen = nDataLen / 16 * 16;
	if ((int)nDataLen == nRainbowChainCount * 16)
	{
		std::cout << "precomputation of this rainbow table already finished" << std::endl;
		fclose(file);
		return 0;
	}
	if (nDataLen > 0)
	{
		std::cout << "continuing from interrupted precomputation..." << std::endl;
	}
	
	//Partfile.seekp(0, std::ios::end);
	size_t nReturn;
	CChainWalkContext cwc;
	uint64 nIndex[2];
	for(m_nCurrentCalculatedChains = nDataLen / 16; m_nCurrentCalculatedChains < nRainbowChainCount; m_nCurrentCalculatedChains++)
	{
		cwc.GenerateRandomIndex();
		nIndex[0] = cwc.GetIndex();
		for (int nPos = 0; nPos < nRainbowChainLen - 1; nPos++)
		{
			cwc.IndexToPlain();
			cwc.PlainToHash();
			cwc.HashToIndex(nPos);
		}

		nIndex[1] = cwc.GetIndex();
		if ((nReturn = fwrite(nIndex, 1, 16, file)) != 16)
		{
			printf("disk write fail\n");
			break;
		}
//		fflush(file);
//
//		nFileLen = (unsigned int)GetFileLen(file);
//		fseek(file, nFileLen, SEEK_SET);
/*		if(nFileLen != ((m_nCurrentCalculatedChains * 16) + 16))
		{
			printf("Invalid file size!");
		}
*/
//		Partfile.flush();
		// Increase the pointer location
	}
	fclose(file);
	/*
	while(nCalculatedChains < nRainbowChainCount)
	{
		for(int i = 0; i < m_nProcessorCount; i++)
		{
			if(m_pThreads[i]->GetIsDataReadyFlag() > 0)
			{

				// Retrieve the calculated data
				const char *data = (const char *)m_pThreads[i]->GetData();
				// Write it to the file
				int nNewChains = DATA_CHUNK_SIZE / 16;
				// Check if too many chains is calculated
				if(nRainbowChainCount < (nCalculatedChains + nNewChains))
				{
					// If this is the fact, reduce the amount of chains we copy
					nNewChains = nRainbowChainCount - nCalculatedChains;
				}
				if(nNewChains == 0)
					break;
				nCalculatedChains += nNewChains;
				
				Partfile.write(data, nNewChains * 16);
				Partfile.flush();

				// And mark the buffer as empty
				m_pThreads[i]->ClearDataReadyFlag();
			}
		}
		Sleep(1);
	}

	// Stop the threads again and destroy them
	for(int i = 0; i < m_nProcessorCount; i++)
	{
		m_pThreads[i]->Stop();
		delete m_pThreads[i];
		m_pThreads[i] = NULL;
	}	

*/
	return 0;
}

